Frog position after T seconds [DFS, Stack, BFS]

Time: O(N); Space: O(N); hard

Given an undirected tree consisting of n vertices numbered from 1 to n. A frog starts jumping from the vertex 1. In one second, the frog jumps from its current vertex to another unvisited vertex if they are directly connected. The frog can not jump back to a visited vertex. In case the frog can jump to several vertices it jumps randomly to one of them with the same probability, otherwise, when the frog can not jump to any unvisited vertex it jumps forever on the same vertex.

The edges of the undirected tree are given in the array edges, where edges[i] = [fromi, toi] means that exists an edge connecting directly the vertices fromi and toi.

Return the probability that after t seconds the frog is on the vertex target.

Example 1:

Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 2, target = 4

Output: 0.16666666666666666

Explanation:

  • The figure above shows the given graph.

  • The frog starts at vertex 1, jumping with 1/3 probability to the vertex 2 after second 1

  • and then jumping with 1/2 probability to vertex 4 after second 2.

  • Thus the probability for the frog is on the vertex 4 after 2 seconds is 1/3 * 1/2 = 1/6 = 0.16666666666666666.

Example 2:

Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 1, target = 7

Output: 0.3333333333333333

Explanation:

  • The figure above shows the given graph.

  • The frog starts at vertex 1, jumping with 1/3 = 0.3333333333333333 probability to the vertex 7 after second 1.

Example 3:

Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 20, target = 6

Output: 0.16666666666666666

Constraints:

  • 1 <= n <= 100

  • len(edges) = n-1

  • len(edges[i]) = 2

  • 1 <= edges[i][0], edges[i][1] <= n

  • 1 <= t <= 50

  • 1 <= target <= n

  • Answers within 10^-5 of the actual value will be accepted as correct.

Hints:

  1. Use a variation of DFS with parameters ‘curent_vertex’ and ‘current_time’.

  2. Update the probability considering to jump to one of the children vertices.

1. BFS solution with better precision

[1]:
import collections

class Solution1(object):
    """
    Time: O(N)
    Space: O(N)
    """
    def frogPosition(self, n, edges, t, target):
        """
        :type n: int
        :type edges: List[List[int]]
        :type t: int
        :type target: int
        :rtype: float
        """
        G = collections.defaultdict(list)
        for u, v in edges:
            G[u].append(v)
            G[v].append(u)

        stk = [(t, 1, 0, 1)]
        while stk:
            new_stk = []
            while stk:
                t, node, parent, choices = stk.pop()
                if not t or not (len(G[node])-(parent != 0)):
                    if node == target:
                        return 1.0/choices
                    continue

                for child in G[node]:
                    if child == parent:
                        continue
                    new_stk.append((t-1, child, node, choices*(len(G[node])-(parent != 0))))
            stk = new_stk

        return 0.0
[2]:
s = Solution1()

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 2
target = 4
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 1
target = 7
assert s.frogPosition(n, edges, t, target) == 0.3333333333333333

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 20
target = 6
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

2. BFS solution with stack and better precision

[3]:
class Solution2(object):
    def frogPosition(self, n, edges, t, target):
        """
        :type n: int
        :type edges: List[List[int]]
        :type t: int
        :type target: int
        :rtype: float
        """
        G = collections.defaultdict(list)
        for u, v in edges:
            G[u].append(v)
            G[v].append(u)

        stk = [(t, 1, 0, 1)]
        while stk:
            t, node, parent, choices = stk.pop()
            if not t or not (len(G[node])-(parent != 0)):
                if node == target:
                    return 1.0/choices
                continue

            for child in G[node]:
                if child == parent:
                    continue
                stk.append((t-1, child, node, choices*(len(G[node])-(parent != 0))))

        return 0.0
[4]:
s = Solution2()

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 2
target = 4
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 1
target = 7
assert s.frogPosition(n, edges, t, target) == 0.3333333333333333

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 20
target = 6
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

3. DFS solution with recursion and better precision

[5]:
import collections

class Solution3(object):
    def frogPosition(self, n, edges, t, target):
        """
        :type n: int
        :type edges: List[List[int]]
        :type t: int
        :type target: int
        :rtype: float
        """
        def dfs(G, target, t, node, parent):
            if not t or not (len(G[node])-(parent != 0)):
                return int(node == target)
            result = 0
            for child in G[node]:
                if child == parent:
                    continue
                result = dfs(G, target, t-1, child, node)
                if result:
                    break
            return result*(len(G[node])-(parent != 0))

        G = collections.defaultdict(list)
        for u, v in edges:
            G[u].append(v)
            G[v].append(u)
        choices = dfs(G, target, t, 1, 0)

        return 1.0/choices if choices else 0.0
[6]:
s = Solution3()

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 2
target = 4
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 1
target = 7
assert s.frogPosition(n, edges, t, target) == 0.3333333333333333

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 20
target = 6
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

4. DFS solution with recursion

[7]:
import collections

class Solution4(object):
    def frogPosition(self, n, edges, t, target):
        """
        :type n: int
        :type edges: List[List[int]]
        :type t: int
        :type target: int
        :rtype: float
        """
        def dfs(G, target, t, node, parent):
            if not t or not (len(G[node])-(parent != 0)):
                return float(node == target)
            for child in G[node]:
                if child == parent:
                    continue
                result = dfs(G, target, t-1, child, node)
                if result:
                    break
            return result/(len(G[node])-(parent != 0))

        G = collections.defaultdict(list)
        for u, v in edges:
            G[u].append(v)
            G[v].append(u)

        return dfs(G, target, t, 1, 0)
[8]:
s = Solution4()

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 2
target = 4
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 1
target = 7
assert s.frogPosition(n, edges, t, target) == 0.3333333333333333

n = 7
edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]
t = 20
target = 6
assert s.frogPosition(n, edges, t, target) == 0.16666666666666666